[TIR][USMP] Added buffer info extraction pass #8468

manupak · 2021-07-14T15:59:19Z

This commit adds a pass that takes the main (call graph of operators)
TIR PrimFunc and each operators also as TIR PrimFunc. The pass will
traverse through all TIR PrimFunc starting the from main. Thereafter,
it will extract information from tir.allocates. Among the information,
the liveness conflicts are reported.

Added test for a linear model
Added test for parallel/serial mixed for loops
Added test for a substructure of inception-style model.
Exposed buffer_info creation to python
Added member functions to update pool info
Unit tests to cover functionality of buffer_info

There is a subsequent PR depends on this : manupak#1

manupak · 2021-08-11T16:08:20Z

I think this could be reviewed now. @areusch @jroesch @jcf94

junrushao · 2021-08-11T18:09:36Z

Hey @manupa-arm, thanks for the nice PR!

As a main thing I am a little bit worried about, I would love to discuss specifically on the IR-related change, because in the worse case, adding/changing/removing an IR node would result in malfunctioning of many passes. Therefore, we tend to keep the IR stable if possible, and scrutinize any changes to the IR. For example, when @masahi was adding the while node (really awesome work #7425), we discussed very carefully on the implication, and actually updated many passes that are affected by that particular change.

As we are adopting the new RFC process, I would love to propose that if the RFC implies any IR modification, the modification should be written explicitly in a section (e.g. which fields are added), and be discussed really carefully on potential alternatives without IR change, and everything that could happen if the IR is changed (e.g. which passes won't work). CC: @jroesch @vinx13 @areusch @comaniac

On this particular case, we are seeing a set of IR changes, including this PR, the PrimFunc PR, the allocate_const PR. That will be super helpful if you guys could send a RFC discussing the rationale and potential implication on this set of changes to TIR, so that it could be more more maintainable in the foreseeable future.

🙏 Many thanks!

manupak · 2021-08-11T19:00:07Z

Hi @junrushao1994 ,

Yes, we are working on writing up the allocate_const RFC and hoping to post it soon. For this one's addition of pinned_memory attribute to allocate, happy to write up a RFC but I feel there is an alternative w/o changing the IR if thats preferred using AttrStmt -- let me know what you think.

tqchen · 2021-08-11T20:56:47Z

Thanks @manupa-arm . I believe @junrushao1994 's comment was mainly about the process.

It would be great if we can separate the IR node change RFC/PRs from the feature RFC(USMP). since IR changes would involve more deliberation, and explicitly marking the RFC about the IR node change would help to get the folks involved in these parts without necessarily going to deep into the feature RFC.

Back to the pinned memory itself, we might be able to fold some of that info into the PointerType storage_scope field as what @masahi did in some of his recent refactor

areusch

@manupa-arm i didn't have time to do a full review but here's some early feedback

include/tvm/tir/stmt.h

include/tvm/tir/usmp/analysis.h

include/tvm/tir/usmp/utils.h

src/tir/usmp/analysis/extract_buffer_info.cc

areusch · 2021-08-19T20:18:29Z

src/tir/usmp/analysis/extract_buffer_info.cc

+      // We only statically memory plan only allocates with known
+      // compile time sizes.
+      auto buffer_info = BufferInfo(op->buffer_var->name_hint, size_bytes);
+      buffer_info->SetPoolCandidates(ParseCommaSeperatedString(op->pinned_memory));


what if pinned_memory is empty?

Explained more and inserted checks

include/tvm/tir/usmp/utils.h

areusch · 2021-09-08T00:48:06Z

i think the current status of this PR is that it's blocked on apache/tvm-rfcs#22 . @manupa-arm do you agree?

manupak · 2021-09-08T06:28:11Z

I think this is blocked on me doing RFC text changes in apache/tvm-rfcs#23 and reflecting that to USMP RFC.

This one does not consider tir.allocate_const / tir.constants just yet; We are planning to add that in a seperate PR once apache/tvm-rfcs#22 is resolved and #8509 is landed.

Sorry, I ll spend the next cycles on apache/tvm-rfcs#23 --> USMP RFC --> this one after landing this : #8795 which keeps getting conflicted

manupak · 2021-10-05T11:08:10Z

This has been updated to reflect the changes as per the discussion on the RFCs.
However, currently it depends on the PR : #9168 , therefore putting this to draft until that is merged.

manupak

@areusch addressed the comments! :)

include/tvm/tir/usmp/utils.h

areusch

@manupa-arm thanks, had a look over this. i think it's mostly straightforward and am mainly requesting some more comments for those less familiar.

cc @mbs-octoml @electriclilies @mikepapadim can you have a look?

include/tvm/tir/usmp/utils.h

python/tvm/tir/usmp/__init__.py

python/tvm/tir/usmp/utils.py

src/tir/usmp/analysis/extract_buffer_info.cc

areusch · 2021-10-12T20:37:44Z

tests/python/unittest/test_tir_usmp_analysis_extract_bufferinfo.py

+import tvm
+from tvm import tir, script
+from tvm.ir import Range
+from tvm.script import tir as T


this might be a bit of too shortened--i think pylint would complain if we ran it on tests

This is following precedence from the code base, all the tvmscript based tests are like this. Thus, better to be consistent. Moreover, I think tvmscript printer uses this form now.

tests/python/unittest/test_tir_usmp_analysis_extract_bufferinfo.py

tests/python/unittest/test_tir_usmp_utils.py

manupak

Thanks @areusch for the review.
PTAL when you have some time.

include/tvm/tir/usmp/utils.h

python/tvm/tir/usmp/__init__.py

python/tvm/tir/usmp/utils.py

src/tir/usmp/analysis/extract_buffer_info.cc

manupak · 2021-10-21T17:39:03Z

tests/python/unittest/test_tir_usmp_analysis_extract_bufferinfo.py

+import tvm
+from tvm import tir, script
+from tvm.ir import Range
+from tvm.script import tir as T


This is following precedence from the code base, all the tvmscript based tests are like this. Thus, better to be consistent. Moreover, I think tvmscript printer uses this form now.

tests/python/unittest/test_tir_usmp_analysis_extract_bufferinfo.py

tests/python/unittest/test_tir_usmp_utils.py

manupak · 2021-11-01T19:33:22Z

a friendly ping @mbs-octoml !

mbs-octoml

Sorry for the delay, finally got to this. LGTM, just some nits and a question about multiple calls to the same PrimFunc with diff buffers.

include/tvm/tir/usmp/utils.h

mbs-octoml · 2021-11-04T22:35:10Z

include/tvm/tir/usmp/utils.h

+  /*! \brief The pool candidates that this buffer can get pooled to*/
+  Array<PoolInfo> pool_candidates;
+  /*! \brief The byte alignment required within the pool */
+  Integer alignment;


nit: I guess this has to be the gcm of the alignment constraints imposed by all user's of buffer -- is there any use for an Array to capture that?

Yes -- that has to be captured at the creation of the tir.allocate node though. Therefore, at the point of BufferInfo creation that is already decided.

python/tvm/tir/usmp/analysis/analysis.py

src/tir/usmp/analysis/extract_buffer_info.cc

mbs-octoml · 2021-11-04T22:50:20Z

src/tir/usmp/analysis/extract_buffer_info.cc

+
+  void UpdateAliases(const Array<PrimExpr>& args, const PrimFunc& func);
+
+  Map<BufferInfo, tir::Stmt> buffer_info_map_;


comments for all these fields please

src/tir/usmp/analysis/extract_buffer_info.cc

mbs-octoml · 2021-11-04T22:57:31Z

src/tir/usmp/analysis/extract_buffer_info.cc

+}
+
+void BufferInfoExtractor::VisitStmt_(const AllocateNode* op) {
+  const auto& currect_scope_info = scope_stack_.top();


nit: current

mbs-octoml · 2021-11-04T23:07:31Z

src/tir/usmp/analysis/extract_buffer_info.cc

+  return buffer_vars;
+}
+
+void BufferInfoExtractor::UpdateAliases(const Array<PrimExpr>& args, const PrimFunc& func) {


At most one call per function, right? Or at least if there are multiple calls then they all use the same buffers. Am I reading that right? CHECK fail if that assumption fails?

this also assumes the calls are seen after the PrimFunc defns, right?

I found this while integrating the USMP :) in my local. Added more code to support per-call analysis that lead sporadic liveness. Also added a test case to show this as well.

Therefore, now this pass could handle multiple calls per PrimFunc.

This commit adds a pass that takes the main (call graph of operators) TIR PrimFunc and each operators also as TIR PrimFunc. The pass will traverse through all TIR PrimFunc starting the from main. Thereafter, it will extract information from tir.allocates. Among the information, the liveness conflicts are reported. * Added test for a linear model * Added test for parallel/serial mixed for loops * Added test for a substructure of inception-style model. * Exposed buffer_info creation to python * Added member functions to update pool info * Unit tests to cover functionality of buffer_info Change-Id: I5e163ac3e83c830629a5d34ed4407c9962701c60

Swap key-value pairs of returned values of the buffer_info extraction pass. Change-Id: Ia4f7289592bc776ef6189a41a7891038751bf31f

Updating the USMP utility tests to include tests that test creation of PoolInfo and PoolAllocation Objects. Change-Id: I5d349d0ffcac6b0160072d832dd9d5418699228e

* Removing the unnecessary header : include/tvm/tir/usmp/analysis.h * Some nits and cleanup Change-Id: Iac3ddd9428c56cd8ef49cf643e797bf6fdf4e97a

* Change the class data members to have a trailing underscore Change-Id: I71809b3c73b0bc0cd133fad1392ae8c17c895ee4

Adding more documentation for data structures and the approach Change-Id: Ide2bfffaeff9add86853b6992017264e5d796299

* Added more documentation * Added functionality to handle multiple calls for the same PrimFunc with a test. Change-Id: Ib7c27b3cf17f415067a224f1e57d8b928f4c7c6f

manupak

@mbs-octoml @areusch I finally got around to address the comments! PTAL when you have some time.

include/tvm/tir/usmp/utils.h

manupak · 2021-11-22T12:03:16Z

include/tvm/tir/usmp/utils.h

+  /*! \brief The pool candidates that this buffer can get pooled to*/
+  Array<PoolInfo> pool_candidates;
+  /*! \brief The byte alignment required within the pool */
+  Integer alignment;


Yes -- that has to be captured at the creation of the tir.allocate node though. Therefore, at the point of BufferInfo creation that is already decided.

python/tvm/tir/usmp/analysis/analysis.py

src/tir/usmp/analysis/extract_buffer_info.cc

manupak · 2021-11-22T12:06:04Z

src/tir/usmp/analysis/extract_buffer_info.cc

+  return buffer_vars;
+}
+
+void BufferInfoExtractor::UpdateAliases(const Array<PrimExpr>& args, const PrimFunc& func) {


I found this while integrating the USMP :) in my local. Added more code to support per-call analysis that lead sporadic liveness. Also added a test case to show this as well.

Therefore, now this pass could handle multiple calls per PrimFunc.

src/tir/usmp/analysis/extract_buffer_info.cc

* Attaching targets to PrimFuncs in the util test case Change-Id: I82960512659a346f6242b2b5789ec1120f8ea2cf

manupak · 2021-11-22T22:07:04Z

@areusch a friendly ping!

areusch · 2021-11-23T17:44:21Z

python/tvm/tir/usmp/utils.py

+            alignment,
+        )
+
+    def set_pool_candidates(self, pool_candidates: list):


note: looks like this is cruft from a previous incarnation of this PR. TODO(@manupa-arm): remove in a follow-on

* [TIR][USMP] Added buffer info extraction pass This commit adds a pass that takes the main (call graph of operators) TIR PrimFunc and each operators also as TIR PrimFunc. The pass will traverse through all TIR PrimFunc starting the from main. Thereafter, it will extract information from tir.allocates. Among the information, the liveness conflicts are reported. * Added test for a linear model * Added test for parallel/serial mixed for loops * Added test for a substructure of inception-style model. * Exposed buffer_info creation to python * Added member functions to update pool info * Unit tests to cover functionality of buffer_info Change-Id: I5e163ac3e83c830629a5d34ed4407c9962701c60 * [TIR][USMP] Added buffer info extraction pass Swap key-value pairs of returned values of the buffer_info extraction pass. Change-Id: Ia4f7289592bc776ef6189a41a7891038751bf31f * [TIR][USMP] Added buffer info extraction pass Updating the USMP utility tests to include tests that test creation of PoolInfo and PoolAllocation Objects. Change-Id: I5d349d0ffcac6b0160072d832dd9d5418699228e * [TIR][USMP] Added buffer info extraction pass * Removing the unnecessary header : include/tvm/tir/usmp/analysis.h * Some nits and cleanup Change-Id: Iac3ddd9428c56cd8ef49cf643e797bf6fdf4e97a * [TIR][USMP] Added buffer info extraction pass * Change the class data members to have a trailing underscore Change-Id: I71809b3c73b0bc0cd133fad1392ae8c17c895ee4 * [TIR][USMP] Added buffer info extraction pass Adding more documentation for data structures and the approach Change-Id: Ide2bfffaeff9add86853b6992017264e5d796299 * [TIR][USMP] Added buffer info extraction pass * Added more documentation * Added functionality to handle multiple calls for the same PrimFunc with a test. Change-Id: Ib7c27b3cf17f415067a224f1e57d8b928f4c7c6f * [TIR][USMP] Added buffer info extraction pass * Attaching targets to PrimFuncs in the util test case Change-Id: I82960512659a346f6242b2b5789ec1120f8ea2cf

manupak mentioned this pull request Jul 14, 2021

[RFC] Unified Static Memory Planning (USMP) Tracking Issue #8404

Closed

21 tasks

manupak force-pushed the usmp_extract_buffer_info branch from f901bdf to a13677d Compare July 19, 2021 14:58

manupak force-pushed the usmp_extract_buffer_info branch from a13677d to 321e30f Compare August 6, 2021 16:57

manupak requested review from areusch, comaniac, jroesch, junrushao, kparzysz-quic, masahi, merrymercy, tqchen, vinx13, yzhliu, zhiics and ZihengJiang as code owners August 6, 2021 16:57

areusch reviewed Aug 19, 2021

View reviewed changes

manupak force-pushed the usmp_extract_buffer_info branch from bf17872 to f006178 Compare October 5, 2021 11:06

manupak requested a review from Hzfengsy as a code owner October 5, 2021 11:06

manupak marked this pull request as draft October 5, 2021 11:06

manupak force-pushed the usmp_extract_buffer_info branch 2 times, most recently from 03d1ba0 to 8294bcf Compare October 6, 2021 18:01

manupak mentioned this pull request Oct 6, 2021

[TIR][USMP] Greedy memory planning algorithm #9214

Merged

manupak force-pushed the usmp_extract_buffer_info branch from 8294bcf to 1954bde Compare October 7, 2021 15:23

manupak commented Oct 7, 2021

View reviewed changes

include/tvm/tir/usmp/utils.h Show resolved Hide resolved

areusch requested changes Oct 12, 2021

View reviewed changes

manupak force-pushed the usmp_extract_buffer_info branch from 3e8a78e to de8e4f1 Compare October 21, 2021 17:30

manupak commented Oct 21, 2021

View reviewed changes

manupak force-pushed the usmp_extract_buffer_info branch from de8e4f1 to 10557aa Compare October 30, 2021 06:40

manupak mentioned this pull request Nov 1, 2021

[TIR][USMP] adding the pass to convert to pool offsets #9418

Merged

mbs-octoml approved these changes Nov 4, 2021

View reviewed changes

manupak added 7 commits November 22, 2021 11:28

[TIR][USMP] Added buffer info extraction pass

60664a2

Swap key-value pairs of returned values of the buffer_info extraction pass. Change-Id: Ia4f7289592bc776ef6189a41a7891038751bf31f

[TIR][USMP] Added buffer info extraction pass

85bb587

Updating the USMP utility tests to include tests that test creation of PoolInfo and PoolAllocation Objects. Change-Id: I5d349d0ffcac6b0160072d832dd9d5418699228e

[TIR][USMP] Added buffer info extraction pass

7cd3bc5

* Removing the unnecessary header : include/tvm/tir/usmp/analysis.h * Some nits and cleanup Change-Id: Iac3ddd9428c56cd8ef49cf643e797bf6fdf4e97a

[TIR][USMP] Added buffer info extraction pass

486facf

* Change the class data members to have a trailing underscore Change-Id: I71809b3c73b0bc0cd133fad1392ae8c17c895ee4

[TIR][USMP] Added buffer info extraction pass

63a155c

Adding more documentation for data structures and the approach Change-Id: Ide2bfffaeff9add86853b6992017264e5d796299

[TIR][USMP] Added buffer info extraction pass

c130133

* Added more documentation * Added functionality to handle multiple calls for the same PrimFunc with a test. Change-Id: Ib7c27b3cf17f415067a224f1e57d8b928f4c7c6f

manupak force-pushed the usmp_extract_buffer_info branch from 10557aa to c130133 Compare November 22, 2021 12:22

manupak requested a review from icemelon as a code owner November 22, 2021 12:22

manupak commented Nov 22, 2021

View reviewed changes

[TIR][USMP] Added buffer info extraction pass

64f1d55

* Attaching targets to PrimFuncs in the util test case Change-Id: I82960512659a346f6242b2b5789ec1120f8ea2cf

manupak mentioned this pull request Nov 23, 2021

[TIR][USMP] Integrating USMP to AoT Executor #9565

Merged

areusch approved these changes Nov 23, 2021

View reviewed changes

areusch merged commit 18730c4 into apache:main Nov 23, 2021

driazati mentioned this pull request Jul 14, 2022

TVM v0.9.0.rc0 Release Candidate Notes #12102

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[TIR][USMP] Added buffer info extraction pass #8468

[TIR][USMP] Added buffer info extraction pass #8468

manupak commented Jul 14, 2021 •

edited

Loading

manupak commented Aug 11, 2021

junrushao commented Aug 11, 2021 •

edited

Loading

manupak commented Aug 11, 2021

tqchen commented Aug 11, 2021 •

edited

Loading

areusch left a comment

areusch Aug 19, 2021

manupak Oct 7, 2021

areusch commented Sep 8, 2021

manupak commented Sep 8, 2021

manupak commented Oct 5, 2021

manupak left a comment

areusch left a comment

areusch Oct 12, 2021

manupak Oct 21, 2021 •

edited

Loading

manupak left a comment

manupak Oct 21, 2021 •

edited

Loading

manupak commented Nov 1, 2021

mbs-octoml left a comment

mbs-octoml Nov 4, 2021

manupak Nov 22, 2021

mbs-octoml Nov 4, 2021

mbs-octoml Nov 4, 2021

mbs-octoml Nov 4, 2021

mbs-octoml Nov 4, 2021

manupak Nov 22, 2021

manupak left a comment

manupak Nov 22, 2021

manupak Nov 22, 2021

manupak commented Nov 22, 2021

areusch Nov 23, 2021


		void UpdateAliases(const Array<PrimExpr>& args, const PrimFunc& func);

		Map<BufferInfo, tir::Stmt> buffer_info_map_;

[TIR][USMP] Added buffer info extraction pass #8468

[TIR][USMP] Added buffer info extraction pass #8468

Conversation

manupak commented Jul 14, 2021 • edited Loading

manupak commented Aug 11, 2021

junrushao commented Aug 11, 2021 • edited Loading

manupak commented Aug 11, 2021

tqchen commented Aug 11, 2021 • edited Loading

areusch left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

areusch commented Sep 8, 2021

manupak commented Sep 8, 2021

manupak commented Oct 5, 2021

manupak left a comment

Choose a reason for hiding this comment

areusch left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

manupak Oct 21, 2021 • edited Loading

Choose a reason for hiding this comment

manupak left a comment

Choose a reason for hiding this comment

manupak Oct 21, 2021 • edited Loading

Choose a reason for hiding this comment

manupak commented Nov 1, 2021

mbs-octoml left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

manupak left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

manupak commented Nov 22, 2021

Choose a reason for hiding this comment

manupak commented Jul 14, 2021 •

edited

Loading

junrushao commented Aug 11, 2021 •

edited

Loading

tqchen commented Aug 11, 2021 •

edited

Loading

manupak Oct 21, 2021 •

edited

Loading

manupak Oct 21, 2021 •

edited

Loading